from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
import json
import requests
from bs4 import BeautifulSoup


def fetch_video_info(url, headers=None):
    # 设置 Chrome options
    chrome_options = Options()
    chrome_options.add_argument("--headless")  # 无头模式
    chrome_options.add_argument("--disable-gpu")

    # 设置 ChromeDriver 的路径
    chrome_service = Service("/usr/local/bin/chromedriver")  # 确保 chromedriver 在正确路径
    # chrome_service = Service("/Users/qiuyuanyu/Downloads/chromedriver-mac-arm64/chromedriver")  # 确保 chromedriver 在正确路径

    # 启动浏览器
    driver = webdriver.Chrome(service=chrome_service, options=chrome_options)

    try:
        # 设置请求头（如果需要）
        if headers:
            for key, value in headers.items():
                driver.execute_cdp_cmd('Network.setExtraHTTPHeaders', {"headers": {key: value}})

        # 打开目标网页
        driver.get(url)

        # 等待 video 标签加载出来之后再执行后续操作
        WebDriverWait(driver, 600).until(EC.presence_of_element_located((By.TAG_NAME, 'video')))
        # 确保 video 的 source 或其他子标签加载完成
        WebDriverWait(driver, 600).until(EC.presence_of_element_located((By.TAG_NAME, "source")))
        # 查找 <video> 标签并获取 src 属性
        video_element = driver.find_element(By.TAG_NAME, 'video')
        # 打印 video 标签 内容
        print(video_element.get_attribute('outerHTML'))
        # 找到 source 标签下 domain 包含 www.douyin.com src 完整属性
        source_elements = video_element.find_elements(By.TAG_NAME, 'source')
        # 如果source_elements 为空，抛出异常结束
        if not source_elements:
            raise Exception("No source elements found")
        video_src = None
        for source in source_elements:
            src = source.get_attribute('src')
            if 'www.douyin.com' in src:
                video_src = src
                break

        if not video_src:
            raise Exception("No valid video source found")

        # webdriver 通过一个新标签 打开  video_src
        driver.execute_script(f"window.open('{video_src}')")
        # 切换到新标签
        driver.switch_to.window(driver.window_handles[-1])
        # 获取页面的 video 标签 的 src 属性
        actual_video_url = driver.find_element(By.TAG_NAME, 'source').get_attribute('src')
        # 构造 JSON 返回值
        result = {
            "aweme": actual_video_url,
            "attribute_2": "additional_attribute_value"  # 可替换为你需要的其他属性
        }
        return result
    finally:
        # 关闭 WebDriver
        driver.quit()
